import numpy as np
import onnxruntime as ort
import sys
from onnx_infer.utils.utils import image_denormalization,image_normalization
import cv2


class Cartoon:
    def __init__(self,model_path) -> None:
        self.ort_session = ort.InferenceSession(model_path)
        self.onnx_input_name = self.ort_session.get_inputs()[0].name
        self.onnx_output_name = self.ort_session.get_outputs()[0].name
    def run(self,image,size):
        height = image.shape[0]
        width = image.shape[1]
        image = cv2.resize(image,[size,size])
        image  = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
        input_data = image_normalization(image)
        input_data = np.expand_dims(input_data,axis=0).astype(np.float32)
        input_data = np.transpose(input_data,(0,3,1,2))
        
        onnx_result = self.ort_session.run([self.onnx_output_name],input_feed={self.onnx_input_name:input_data})[0]
        
        output_data = np.transpose(onnx_result,(0,2,3,1))[0]
        
        output_data = cv2.cvtColor(image_denormalization(output_data),cv2.COLOR_RGB2BGR)
        return cv2.resize(output_data,(width,height))
    
    
if __name__ == "__main__":
    model_path = r'cartoon\weight\best_3d_weight\generator_best_512.onnx'
    
    model = Cartoon(model_path)
    input_filename = r'asset\1.jpg'
    img = cv2.imread(input_filename)
    
    img = model.run(img,512)
    cv2.imwrite('lll.jpg',img)
    # cv2.imshow('dsa',img)
    # cv2.waitKey()
    